# 根据文章摘要生成插图
"""
最后一个就是 ImageNode，基于文章的摘要配插图。图片生成需要多模态大模型，这里我们使用豆包的

seedream-3.0-t2i 这个文生图模型，它的图片生成质量还是不错的。（PS. 豆包大模型可以通过火山引擎来使用，新注册的账号还有一定的免费额度）。

生图成功后，会返回一个可访问的 url ，但是这个 url 通常是有过期时间的，所以我把图片 download 下来保存到本地，方便后面合并到文章中。
"""

import os
import dotenv
import requests
from openai import OpenAI
from Article_image_ai_agent.state import ArticleState

# 生图Prompt
IMAGE_PROMPT = """
你是一位资深的插画师，擅长为各种文章创作精美的插图。
现在有一篇文章，具体信息如下：
主题：{topic}
标题：{title}
摘要：{summary}
请你为该文章创造一幅插画，展示文章所要传达的核心内容。要求画质精美、表达力强，且意境深远，给人无尽回味。
特别注意：图片中的文字全部采用英文。
"""

# 加载环境变量
dotenv.load_dotenv()

# 创建豆包Seedream客户端
seedream = OpenAI(
    base_url=os.getenv("VOLCENGINE_API_BASE"),
    api_key=os.environ.get("VOLCENGINE_API_KEY"),
)


def image_node(state: ArticleState) -> ArticleState:
    """
    文章插图节点
    :param state: 当前文章状态
    :return: 处理后的文章状态
    """

    # 从当前状态中获取所需信息
    if not state["summary"]:
        raise ValueError("未生成文章摘要！")
    summary = state["summary"]

    # 使用豆包的Seedream模型生成图片
    response = seedream.images.generate(
        model="doubao-seedream-3-0-t2i-250415",
        prompt=summary,
        size="1024x1024",
        response_format="url"
    )

    if len(response.data) == 0:
        raise ValueError("生成图片失败！")
    print(f"文章插图生成成功！url: {response.data[0].url}")

    # 从url下载图片，并以png格式保存到本地
    image_path = f"./{state["title"]}.png"
    download_image(response.data[0].url, image_path)

    # 更新状态
    state["image_path"] = image_path
    print(f"文章插图生成完成，已保存至：{image_path} ")

    # 返回更新后的状态
    return state


def download_image(url: str, file_path: str) -> None:
    """
    下载图片
    :param url: 图片url
    :param file_path: 图片保存路径
    :return: None
    """

    # 下载图片
    response = requests.get(url)
    with open(file_path, "wb") as f:
        f.write(response.content)